"""
使用大模型多标签分类
"""

import os
from datasets import load_dataset, load_from_disk, concatenate_datasets

from prompt import get_industry_trans_func
from utils import glm4_vllm, load_obj
from setting import StaticValues
from parse import reason_label_parse

call_vllm = None


def run_multi_cls(name, dataset, call_vllm):
    """
    dataset 是某个企业的全部数据集
    """
    sv = StaticValues(name)
    logger = sv.logger

    attr_name = "multi_cls_prompt"

    if not os.path.exists(sv.INDUSTRY_VLLM_MULTI_OBJ):
        new_dataset = dataset.map(
            get_industry_trans_func(attr_name, sv.multi_cls_prompt),
        )
        prompts = []
        for item in new_dataset:
            prompts.append([{"role": "user", "content": item[attr_name]}])
        # glm4_vllm(prompts, output_dir=sv.INDUSTRY_VLLM_MULTI_OBJ)
        call_vllm(prompts, output_dir=sv.INDUSTRY_VLLM_MULTI_OBJ)
    else:
        logger.info(f"{sv.INDUSTRY_VLLM_MULTI_OBJ} already exists.")

    vllm_output = load_obj(sv.INDUSTRY_VLLM_MULTI_OBJ)

    reasons = []
    labels = []

    for _, pred in enumerate(vllm_output):
        pred_text = pred.outputs[0].text
        reason, label = reason_label_parse(pred_text)
        reasons.append(reason)
        labels.append(label)

    if "reason" in dataset.column_names:
        dataset = dataset.remove_columns(["reason"])
    if "label" in dataset.column_names:
        dataset = dataset.remove_columns(["label"])

    tmp_dataset = dataset.add_column("reason", reasons)
    tmp_dataset = tmp_dataset.add_column("label", labels)
    # tmp_dataset.to_csv(sv.INDUSTRY_CSV, index=False)
    return tmp_dataset


def trans_multi_cls_dataset(name):
    """
        转换LLM预测的label字符串到多标签（one-hot形式）
        原本设计成 one-hot形式是为了方便训练出 多label。但这个太复杂了，后续也没有考虑这样做。
        后续还是会将其转换为 0,1,2...等形式
    """
    sv = StaticValues(name=name)
    LABEL_NAME = sv.LABEL_NAME

    ## 加载大模型的预测结果
    def _parse_func(item):
        # 针对 label 进行调整
        label_value = item["label"]
        for col in LABEL_NAME:
            item[col] = 0
            if col in label_value:
                item[col] = 1
        return item

    return _parse_func


# def preprocess_data(examples):
#     # take a batch of texts
#     text = examples["Tweet"]
#     # encode them
#     # encoding = tokenizer(text, padding="max_length", truncation=True, max_length=128)
#     encoding = tokenizer(text, truncation=True, max_length=128)
#     # add labels
#     labels_batch = {k: examples[k] for k in examples.keys() if k in labels}
#     # create numpy array of shape (batch_size, num_labels)
#     labels_matrix = np.zeros((len(text), len(labels)))
#     # fill numpy array
#     for idx, label in enumerate(labels):
#         labels_matrix[:, idx] = labels_batch[label]
#     encoding["labels"] = labels_matrix.tolist()
#     return encoding


if __name__ == "__main__":
    # run_industry_cls("hydrogen")
    pass
